#include "SysLib.h"

#include "Protocal.h"
#include "Common.h"


PClientInfo CONN_LIST[256];
int CONN_COUNT = 0;

int RLIST[256];
int RCOUNT = 0;

int LAST_SWITCH_PORT_TIME = 0;

PConfig CONFIG;

struct sockaddr_in *SERVER_ADDR;
int LISTEN_SOCKET;



void remove_rlist(int sock_fd) {
    int i;
    for (i=0; i<RCOUNT; i++) {
        if (RLIST[i] == sock_fd) {
            RCOUNT -= 1;
            break;
        }
    }
    for (; i<RCOUNT; i++) {
        RLIST[i] = RLIST[i+1];
    }
}


void close_client_conn(int i) {
    PClientInfo conn = CONN_LIST[i];
    remove_rlist(conn->sock_write);
    close(conn->sock_write);
    if (conn->sock_write2) {
        remove_rlist(conn->sock_write2);
        close(conn->sock_write2);
    }
    CONN_COUNT--;
    for (; i<CONN_COUNT; i++) {
        CONN_LIST[i] = CONN_LIST[i+1];
    }
}


PClientInfo get_conn_by_src_addr(struct sockaddr_in *addr) {
    int i;
    for (i=0; i<RCOUNT; i++) {
        PClientInfo conn = CONN_LIST[i];
        if (conn->src_ip == addr->sin_addr.s_addr && conn->src_port == addr->sin_port) {
            return conn;
        }
    }
    return NULL;
}


PClientInfo get_conn_by_sock_write(int sock_write) {
    int i;
    for (i=0; i<RCOUNT; i++) {
        PClientInfo conn = CONN_LIST[i];
        if (conn->sock_write == sock_write && conn->sock_write2 == sock_write) {
            return conn;
        }
    }
    return NULL;
}


void check_pack_id(PClientInfo conn, int pid) {
    if (!conn) {
        return;
    }
    int last_packet_id = conn->last_packet_id;
    if (conn->last_packet_id >= 9000 && pid <= 1000) {
        last_packet_id -= 10000;
    }
    int diff = pid - last_packet_id;
    if (diff > 1) {
        log_debug("Packet disorder: %d, diff: %d", pid, diff);
        conn->lost_packets += diff - 1;
    } else {
        conn->recv_packets += 1;
    }
    conn->last_packet_id = pid;
}

int next_pack_id(PClientInfo conn) {
    if (conn) {
        return conn->next_packet_id = (conn->next_packet_id + 1) % 10000;
    }
    return 0;
}


void send_control_pack(PClientInfo conn, int sock_write, int sin_port,
        const char *format, ...) {
    static char raw_data[256];
    static char encoded_data[2000];
    va_list argp;
    va_start(argp, format);
    vsprintf(raw_data, format, argp);
    va_end(argp);
    DecodedPacket packet;
    struct sockaddr_in remote_addr;
    packet.pack_id = next_pack_id(conn);
    packet.flag = FLAG_CONTROL;
    packet.data = raw_data;
    packet.size = strlen(packet.data);
    int res_size = 0;
    int result = encode_data(&packet, encoded_data, &res_size);
    if (result == 0) {
        remote_addr.sin_family = AF_INET;
        remote_addr.sin_addr.s_addr = conn->remote_ip;
        remote_addr.sin_port = sin_port;
        sendto(sock_write, encoded_data, res_size, 0,
                (struct sockaddr *)&remote_addr, sizeof(remote_addr));
    }
}

void check_timeout() {
    int i = 0;
    int now = time(0);
    while (i < CONN_COUNT) {
        PClientInfo conn = CONN_LIST[i];
        if (now - conn->last_send_time > CONFIG->timeout
                || now - conn->last_recv_time > CONFIG->timeout) {
            log_info("[Timeout] id: %d, src_port: %r",
                    conn->conn_id, conn->src_port);
            close_client_conn(i);
        } else {
            i += 1;
        }
    }
}


void do_cron_task() {
    static int last_check_timeout = 0;
    static int last_report_time = 0;
    int now = time(0);
    if (now - last_check_timeout > CONFIG->timeout) {
        last_check_timeout = now;
        int i;
        for (i = CONN_COUNT - 1; i >= 0; i--) {
            PClientInfo conn = CONN_LIST[i];
            if (now - conn->last_send_time > CONFIG->timeout
                    || now - conn->last_recv_time > CONFIG->timeout) {
                log_info("[Timeout] id: %d", conn->conn_id);
                close_client_conn(i);
            }
        }
    }

    if (now - last_report_time > CONFIG->report_interval) {
        last_report_time = now;
        int i;
        for (i = 0; i < CONN_COUNT; i++) {
            PClientInfo conn = CONN_LIST[i];
            send_control_pack(conn, conn->sock_write, conn->remote_port,
                    "%s\nreport\n%d\n%d", CONFIG->user,
                    conn->lost_packets, conn->recv_packets);
            conn->lost_packets = 0;
            conn->recv_packets = 0;
        }
    }
}

char * get_control_info(char *data_head) {
    char *head = strchr(data_head, '\n');
    if (head) {
        head[0] = '\0';
        return head + 1;
    } else {
        return NULL;
    }
}

int handle_control_packet(PClientInfo conn, int sock_write,
        struct sockaddr_in *src_addr, PDecodedPacket packet) {
    char *cmd = get_control_info(packet->data);
    if (cmd) {
        log_trace("Recv control cmd: %s", cmd);
        char *arg1 = get_control_info(cmd);
        if (0 == strcmp(cmd, "switch")) {
            if (conn->sock_write != sock_write) return 0;

            int port = str2int(arg1);
            char *arg2 = get_control_info(arg1);
            get_control_info(arg2);
            conn->conn_id = str2int(arg2);
            log_debug("Switch to port %d, id: %d", port, conn->conn_id);
            int sin_port = htons(port);
            if (SERVER_ADDR->sin_port != sin_port) {
                SERVER_ADDR->sin_port = sin_port;
                LAST_SWITCH_PORT_TIME = time(0);
            }

            if (conn->remote_port2 != sin_port) {
                conn->remote_port2 = sin_port;
                if (conn->sock_write2 > 0) {
                    close(conn->sock_write2);
                    remove_rlist(conn->sock_write2);
                }
                conn->sock_write2 = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
                if (conn->sock_write2 > 0) {
                    RLIST[RCOUNT++] = conn->sock_write2;
                }
            }
            if (conn->sock_write2 > 0) {
                send_control_pack(conn, conn->sock_write2, conn->remote_port2,
                        "%s\nrenew\n%d", CONFIG->user, conn->conn_id);
            }
        } else if (0 == strcmp(cmd, "ok")) {
            if (src_addr->sin_port == SERVER_ADDR->sin_port
                    && src_addr->sin_port == conn->remote_port2) {
                conn->sock_write2 = conn->sock_write;
                conn->sock_write = sock_write;
                conn->remote_port = conn->remote_port2;
                conn->remote_port2 = 0;
                send_control_pack(conn, sock_write, conn->remote_port,
                        "%s\nok", CONFIG->user);
                log_info("[Renew] id: %d", conn->conn_id);
            }
        }
    }
    return 0;
}


int handle_sock_write_recv(PClientInfo conn, int sock_write,
        char *buffer, PDecodedPacket packet) {
    struct sockaddr_in addr;
    socklen_t addrlen = sizeof(addr);

    int recv_size = recvfrom(sock_write, buffer, BUFFER_SIZE, 0,
            (struct sockaddr *)&addr, &addrlen);
    if (addr.sin_addr.s_addr != SERVER_ADDR->sin_addr.s_addr) return 1;
    if (recv_size <= 0) return 1;
    buffer[recv_size] = 0;
    if (0 != decode_data(buffer, recv_size, packet)) return 1;
    check_pack_id(conn, packet->pack_id);
    conn->last_recv_time = time(0);
    if (packet->flag & FLAG_CONTROL) {
        return handle_control_packet(conn, sock_write,
                &addr, packet);
    } else {
        addr.sin_family = AF_INET;
        addr.sin_port = conn->src_port;
        addr.sin_addr.s_addr = conn->src_ip;
        sendto(conn->sock_src, packet->data, packet->size, 0,
                (struct sockaddr *)&addr, addrlen);
        return 0;
    }
}


int update_server_port() {
    int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (fd == INVALID_SOCKET) {
        log_error("socket() return error");
        return 1;
    }
    log_debug("New socket");
//    int flag = 1;
//    if (-1 == setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(flag))) {
//        log_error("Error setting tcp no delay");
//        close(fd);
//        return 1;
//    }
//    log_debug("Set sockopt");
    struct sockaddr_in *addr = sockaddr_by_ip_port(CONFIG->server_ip, 80);
    log_debug("Connecting %s:%d", CONFIG->server_ip, 80);
    if (SOCKET_ERROR == connect(fd, (struct sockaddr *)addr, sizeof(struct sockaddr_in))) {
        log_error("Error connect, errno: %d", errno);
        close(fd);
        return 1;
    }

    log_debug("Connected");

    log_debug("Sending sync req");
    char req[2000];
    sprintf(req, "GET %s HTTP/1.1\nHost: public.hzbn.tk\nAccept: */*\nConnection: close\n\n", CONFIG->sync_url);
    send(fd, req, strlen(req), 0);
    log_debug("Sync req sent");
    char buffer[1024];

    int len = 0;
    while (true) {
        int rcount = recv(fd, buffer + len, 1024 - len, 0);
        log_debug("Read %d bytes", rcount);
        if (rcount <= 0) break;
        len += rcount;
    }
    buffer[len] = '\0';
    char *i1 = strstr(buffer, "\n\n");
    char *i2 = strstr(buffer, "\r\n\r\n");
    int index = 2000;
    if (i1) index = i1 - buffer + 2;
    if (i2) index = min(index, i2 - buffer + 4);

    DecodedPacket packet;
    packet.data = malloc(100 * sizeof(char));
    decode_data(buffer + index, len - index, &packet);
    packet.data[packet.size] = '\0';
    int server_port = str2int(packet.data);
    log_info("Sync server port: %d", server_port);
    SERVER_ADDR->sin_port = htons(server_port);
    free(packet.data);
    close(fd);

    return 0;
}

int main(int argc, char **argv) {
    {

        CONFIG = load_config(argc, argv);
        SERVER_ADDR = sockaddr_by_ip_port(CONFIG->server_ip, 80);
        struct sockaddr_in *bind_addr = sockaddr_by_ip_port(CONFIG->bind_ip, CONFIG->bind_port);
        LISTEN_SOCKET = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
        if (0 != bind(LISTEN_SOCKET, (struct sockaddr *)bind_addr, sizeof(struct sockaddr_in))) {
            log_error("Listen %s:%d fail", CONFIG->bind_ip, CONFIG->bind_port);
            exit(1);
        }
        free(bind_addr);

        log_info("Listen %s:%d", CONFIG->bind_ip, CONFIG->bind_port);
        update_server_port();
    }
    fd_set rd_set;
    struct timeval timeout;

    char *buffer = (char *)malloc(BUFFER_SIZE * sizeof(unsigned char));
    PDecodedPacket packet = NULL;
    zmalloc(&packet, sizeof(DecodedPacket));
    packet->data = (char *)malloc(BUFFER_SIZE * sizeof(unsigned char));

    while (true) {
        int i;
        do_cron_task();

        FD_ZERO(&rd_set);

        FD_SET(LISTEN_SOCKET, &rd_set);
        int max_fd = LISTEN_SOCKET;
        for (i = 0; i < CONN_COUNT; i++) {
            PClientInfo conn = CONN_LIST[i];
            FD_SET(conn->sock_write, &rd_set);
            max_fd = max(max_fd, conn->sock_write);
            if (conn->sock_write2) {
                FD_SET(conn->sock_write2, &rd_set);
                max_fd = max(max_fd, conn->sock_write2);
            }
        }

        timeout.tv_sec = 1;
        timeout.tv_usec = 0;
        int rcount = select(max_fd + 1, &rd_set, NULL, NULL, &timeout);
        log_debug("Select %d", rcount);
        if (rcount <= 0) continue;
        if (FD_ISSET(LISTEN_SOCKET, &rd_set)) {
            struct sockaddr_in src_addr;
            socklen_t addr_len = sizeof(src_addr);
            log_debug("Recving from listen socket");
            int recv_size = recvfrom(LISTEN_SOCKET, packet->data, BUFFER_SIZE, 0,
                    (struct sockaddr *)&src_addr, &addr_len);
            log_debug("Recv %d byte", recv_size);
            if (recv_size > 0) {
                packet->data[recv_size] = 0;
                PClientInfo conn = get_conn_by_src_addr(&src_addr);
                if (!conn) {
                    conn = new_client_info(&src_addr, LISTEN_SOCKET, SERVER_ADDR);
                    time_t now = time(0);
                    if (now - LAST_SWITCH_PORT_TIME > CONFIG->renew_interval) {
                        update_server_port();
                    }
                    CONN_LIST[CONN_COUNT++] = conn;
                    RLIST[RCOUNT++] = conn->sock_write;
                    log_info("[Accept] port: %d", src_addr.sin_port);
                    send_control_pack(conn, conn->sock_write, conn->remote_port,
                            "%s\nwhoami", CONFIG->user);
                }
                packet->pack_id = next_pack_id(conn);
                packet->flag = 0;
                packet->size = recv_size;
                int res_size = 0;
                int result = encode_data(packet, buffer, &res_size);
                if (result == 0) {
                    struct sockaddr_in dst_addr;
                    dst_addr.sin_family = AF_INET;
                    dst_addr.sin_port = conn->remote_port;
                    dst_addr.sin_addr.s_addr = conn->remote_ip;
                    sendto(conn->sock_write, buffer, res_size, 0,
                            (struct sockaddr *)&dst_addr, sizeof(dst_addr));
                    conn->last_send_time = time(0);
                }
            }
        }

        for (i=0; i<CONN_COUNT; i++) {
            PClientInfo conn = CONN_LIST[i];
            if (FD_ISSET(conn->sock_write, &rd_set)) {
                handle_sock_write_recv(conn, conn->sock_write, buffer, packet);
            }

            if (conn->sock_write2 && FD_ISSET(conn->sock_write2, &rd_set)) {
                handle_sock_write_recv(conn, conn->sock_write2, buffer, packet);
            }
        }
    }
}
